#!/usr/bin/env python
from rich.markdown import Markdown
from rich.console import Console
from prompt_toolkit import prompt
from prompt_toolkit.key_binding import KeyBindings
import json
import readline
import os
import requests
from openai import OpenAI
import openai
api_key='ms-818639f6-cb9c-4e55-a1cf-1e67d11b6833'
console = Console()
kb = KeyBindings()

@kb.add('enter')  # 相当于 Alt+Enter
def _(event):
    event.current_buffer.validate_and_handle()


@kb.add('tab')
def _(event):
    event.current_buffer.insert_text('\n')


fileuploaded = False



class Base:
    @staticmethod
    def save_config(filename, jsonc):
        with open(filename, "w") as f:
            json.dump(jsonc, f)
    @staticmethod
    def LoadJson(json, jsonf):
        while True:
            print("请选择要修改的设置")
# 使用 enumerate 来为每一行添加序号
            for index, (key, value) in enumerate(json.items(), start=1):
                print(f"{index}.{key}:{value}")
            print(str(len(json)+1)+".保存并退出设置")
            try:
                setnum = int(input())
                for index, (key, value) in enumerate(json.items(), start=1):
                    if setnum == index:
                        setvalue=input("请输入"+key+"的值\n")
                        json[key]=setvalue
                    if setnum==len(json)+1:
                        Base.save_config(json, jsonf) 
                        break
            except:
                print("退出设置")
                break


class Set:
    @staticmethod
    def load_config():
        if os.path.exists("config.json"):
            with open("config.json", "r") as f:
                return json.load(f)

    CONFIG_FILE = "config.json"
    EnableSearch = False
    EnableGeneratePic = False
    EnableDeepThink = False
    functionnum = "000"
    config = load_config()
    @staticmethod
    def SetDialog():
        Base.LoadJson(Set.config, "config.json")
    @staticmethod
    def calc_functionnum():
        Set.functionnum = str(int(Set.EnableSearch)) + str(int(Set.EnableDeepThink)) + str(int(Set.EnableGeneratePic))


class ModelManager:
    @staticmethod
    def ShowModelSetDialog():
        while True:
            print("输入数字来选择要设置的参数")
            print("1.tempreture")


class TextProcesser:
    message = []

    @staticmethod
    def clean_markdown(text):
        return text

    @staticmethod
    def GetLength(text):
        length = 0
        for content in text:
            length += len(content["content"])
        return length

    @staticmethod
    def CheckLen(text):
        while TextProcesser.GetLength(text) > 100000:
            del text[0]
        return text


class AgentAI:
    def GetDeepAnswer():
        pass
    @staticmethod
    def GetAnswer(filelist, message): 
        url = "https://dashscope.aliyuncs.com/api/v1/apps/f89d44230ea5413a91eef2512db83271/completion"
        headers = {'Authorization': "Bearer sk-127dd2cb8e454d16b58f376d820eec63", 'Content-Type': "application/json","X-DashScope-SSE": "enable"}
        body = {
        "input": {
            "prompt": message,
            "enable_web_search": True
        },
        "parameters": {
            "incremental_output": True,
            "enable_web_search": True,
            "rag_options": {
                "session_file_ids": id
            }
        },
        "debug": {"enable_web_search": True},
        "enable_web_search": True
            }
        with requests.post(url, headers=headers, json=body, stream=True) as response:
            full_response = ""
            for line in response.iter_lines():
                if line:                        # 解码并处理每一行
                    decoded_line = line.decode('utf-8')
                    if decoded_line.startswith('data:'):
                        try:
                            json_data = json.loads(decoded_line[5:])
                            chunk_text = json_data['output']['text']
                            full_response += chunk_text
                            print(chunk_text, end='', flush=True)
                            full_response += chunk_text
                            if json_data['output']['finish_reason'] == 'stop':
                                break
                        except json.JSONDecodeError:
                            continue
        return full_response




class PublicAI:
    @staticmethod
    def GetImage(prompt):
        # 获取 API Key
        api_key = "sk-127dd2cb8e454d16b58f376d820eec63"

        # 定义请求 URL
        url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation"

        # 定义请求头
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_key}"
        }

        # 定义请求体
        data = {
            "model": "qwen-image-plus",
            "input": {
                "messages": [
                    {
                        "role": "user",
                        "content": [
                            {
                                "text": prompt
                            }
                        ]
                    }
                ]
            },
            "parameters": {
                "negative_prompt": "",
                "prompt_extend": True,
                "watermark": True,
                "size": "1328*1328"
            }
        }

        # 发送 POST 请求
        response = requests.post(url, headers=headers, json=data)

        # 输出响应结果
        if response.status_code == 200:
            print("请求成功！响应内容：")
            print(response.json())
        else:
            print(f"请求失败，状态码：{response.status_code}，错误信息：{response.text}")
    @staticmethod
    def GetDeepWebAnswer(message):
        api_key = "Bearer sk-127dd2cb8e454d16b58f376d820eec63"
        url = "https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions"
        headers = {'Authorization': api_key, 'Content-Type': "application/json"}
        body = {
            "model": "qwen-plus-latest",
            "presence_penalty":2.0,
            "temperature":0.3,
            "messages": message,
            "enable_search": True,
            "search_options": {
            "forced_search": True,
            "enable_search_extension": True,
            "search_strategy": "max"
    },
 
            "stream": True,
            "stream_options": {"include_usage": True},
            "enable_thinking": True
        }
        full_response = ""
        isFirstContent = True
        try:
            response = requests.post(url=url, json=body, headers=headers, stream=True)
            response.raise_for_status()
            for chunk in response.iter_lines():
                if chunk and b'[DONE]' not in chunk:
                    if chunk.startswith(b'data:'):
                        data_str = chunk[5:].strip().decode('utf-8')
                    else:
                        continue  
                    try:
                        data = json.loads(data_str)
                        choices = data.get('choices', [])
                        if not choices:
                            continue
                        delta = choices[0].get('delta', {})
                        # 处理思维链内容（直接清理后输出）
                        if delta.get('reasoning_content'):
                            cleaned_reasoning = TextProcesser.clean_markdown(delta['reasoning_content'])
                            print(cleaned_reasoning, end="", flush=True)
                        # 处理正式回复内容
                        content = delta.get('content')
                        if content:
                            cleaned_content = TextProcesser.clean_markdown(content)
                            if isFirstContent:
                                print("\n以上为思维链内容，模型回复内容如下\n")
                                isFirstContent = False
                            print(cleaned_content, end="", flush=True)
                            full_response += cleaned_content
                    except (json.JSONDecodeError, KeyError, IndexError):
                        continue
        except requests.exceptions.RequestException as e:
            print(f"\nAPI请求失败: {str(e)}")
            return "请求失败，请检查网络和API密钥"
        return full_response

    @staticmethod
    def GetDeepAnswer(message):
        client = OpenAI(
        base_url = 'https://api-inference.modelscope.cn/v1',
        api_key = 'ms-818639f6-cb9c-4e55-a1cf-1e67d11b6833',   # ModelScope Token
    ) 
        response = client.chat.completions.create(
           # model='Qwen/Qwen3-235B-A22B-Instruct-2507', # ModelScope Model-Id
            model='MiniMax/MiniMax-M2',
            messages=message,
            extra_body={"enable_thinking": False},
            stream=True
        )
        full_response = ""
        for chunk in response:
            print(chunk.choices[0].delta.content, end='', flush=True)
            full_response += str(chunk.choices[0].delta.content)
        print("")
        console.print(Markdown(full_response))
        return full_response
    @staticmethod
    def GetAnswer(message):
        client = OpenAI(
        base_url='https://api-inference.modelscope.cn/v1',
        api_key='ms-818639f6-cb9c-4e55-a1cf-1e67d11b6833',   #  ModelScope Token
    ) 
        response = client.chat.completions.create(
            model='Qwen/Qwen3-235B-A22B-Instruct-2507', # ModelScope Model-Id
            messages=message,
            extra_body={"enable_thinking": False},
            stream=True
        )
        full_response = ""
        for chunk in response:
            print(chunk.choices[0].delta.content, end='', flush=True)
            full_response += str(chunk.choices[0].delta.content)
        print("")
        console.print(Markdown(full_response))
        return full_response
    
    @staticmethod 
    def GetWebAnswer(message):
        api_key = "Bearer sk-127dd2cb8e454d16b58f376d820eec63"
        headers = {'Authorization': api_key, 'Content-Type': "application/json"}
        body = {
            "model": "qwen-plus-2025-04-28",
            "messages": message,
            "enable_search": True,
            "search_options": {
            "forced_search": False,
             "enable_search_extension":True,
             "search_strategy": "max"
    },
            "stream": True,
            "stream_options": {"include_usage": True},
            "enable_thinking": False
            }
        url = "https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions"
        full_response = ""
        try:
            response = requests.post(url=url, json=body, headers=headers, stream=True)
            response.raise_for_status()      
            for chunk in response.iter_lines():
                if not chunk or chunk.strip() == b'data: [DONE]':
                    continue          
                if chunk.startswith(b'data: '):
                    data_org = chunk[6:]
                    try:
                        data = json.loads(data_org)
                        if 'choices' in data and len(data['choices']) > 0:
                            delta = data['choices'][0].get('delta', {})
                            if 'content' in delta:
                                content = delta['content']
                                if content!="None":
                                    # 清理Markdown后输
                                    print(content, end="", flush=True)
                                    full_response += str(content)
                    except json.JSONDecodeError:
                        continue
        except requests.exceptions.RequestException as e:
            print(f"\n请求发生错误: {e}")
            return f"请求发生错误: {e}"
        return full_response

class HistoryManager:
    ChatHistory = []
    IsNewChat = True
    IsChanged = False

    @staticmethod
    def GetText(text, role, content):
        jsoncon = {"role": role, "content": content}
        text.append(jsoncon)
        return text

    @staticmethod
    def LoadHistory():
        history_dir = HistoryManager.GetHistoryDir()
        if not os.path.exists(history_dir):
            print(f"错误：目录 {history_dir} 不存在")
            return
        file_list = [f for f in os.listdir(history_dir) if os.path.isfile(os.path.join(history_dir, f))]
        file_list.sort()
        for idx, filename in enumerate(file_list, 1):
            print(f"{idx}.{filename}")
        if len(file_list) < 1:
            print("还没有历史记录")
            return
        while True:
            try:
                selected = int(input("\n请输入文件序号，输入任意非数字退出"))
                if 1 <= selected <= len(file_list):
                    selected_file = file_list[selected-1]
                    full_path = os.path.join(history_dir, selected_file)
                    print(f"\n已选择文件：{selected_file}")
                    print(f"完整路径：{full_path}")
                    print("请选择操作：1.打开  2.删除  其余任意按键退出")
                    option = input()
                    file_path = os.path.join(history_dir, selected_file)
                    if option == "1":
                        if not os.path.exists(file_path):
                            print("本历史记录已被删除")
                            HistoryManager.LoadHistory()
                        with open(file_path, 'r', encoding='utf-8') as f:
                            HistoryManager.ChatHistory = json.load(f)
                            HistoryManager.CurrentFile = selected_file
                            HistoryManager.IsNewChat = False
                            break
                    elif option == "2":
                        os.remove(file_path)
                        print("文件已删除")
                        HistoryManager.LoadHistory()
                else:
                    print("错误：序号超出范围")
            except ValueError:
                print("退出历史记录操作")
                break

    @staticmethod
    def GetHistoryDir():
        current_dir = os.path.dirname(os.path.abspath(__file__))
        history_path = os.path.join(current_dir, 'history')
        os.makedirs(history_path, exist_ok=True)
        return history_path

    @staticmethod
    def SaveHistory(filename):
        history_dir = HistoryManager.GetHistoryDir()
        txtfolder=os.path.join(history_dir,"txt")
        txtpath=os.path.join(txtfolder,filename)
        with open(txtpath,"w",encoding="utf-8")as f:
            for i,msg in enumerate(HistoryManager.ChatHistory,1):
                f.write(f"--- 第 {i} 轮对话 ---\n")
                f.write(f"角色: {msg['role']}\n")
                f.write("内容:\n")
                f.write(msg["content"])
                f.write("\n\n")  # 段落之间空两行，提高可读性

        file_path = os.path.join(history_dir, filename)
        with open(file_path, 'w', encoding='utf-8') as f:
            json.dump(HistoryManager.ChatHistory, f, ensure_ascii=False)

    @staticmethod
    def SaveDialog():
        if HistoryManager.IsChanged:
            IsSave = input("您刚才进行了对话，是否保存？输入y保存：")
            if IsSave == "y": 
                if HistoryManager.IsNewChat:
                    SaveFileName = input("请输入文件名（不输入将放弃保存）：")
                    if SaveFileName:
                        HistoryManager.SaveHistory(SaveFileName)
                else:
                    HistoryManager.SaveHistory(HistoryManager.CurrentFile)
            HistoryManager.IsChanged = False
    @staticmethod       
    def PrintHistory(): 
        for i,msg in enumerate(HistoryManager.ChatHistory,1):
            print(f"--- 第 {i} 轮对话 ---\n")
            print(f"角色: {msg['role']}\n")
            print("内容:\n")
            print(msg["content"])
            print("\n\n")  # 段落之间空两行，提高可读性
class Restart():
    @staticmethod
    def Restart():
        try:
            with open("ai.config","r")as f:
                sta = f.readline()
                if sta == "true":
                    print("上次未正确退出，输入y恢复历史记录")
                    usercmd = input()
                    if usercmd == "y":
                        with open("/home/shell/history/tmp", 'r', encoding='utf-8') as f:
                            HistoryManager.ChatHistory = json.load(f)
                            HistoryManager.IsNewChat = True
                    else:
                        os.remove("/home/shell/history/tmp")
        except:
            pass
    @staticmethod
    def WriteTmp():
        HistoryManager.SaveHistory("tmp") 

if __name__ == "__main__":
    Restart.Restart()
    os.system("rm ai.config")
    with open("ai.config","w")as f:
        f.write("true")
    print("欢迎使用畅言平板AI系统，输入help查看使用说明")
    while True:
        Restart.WriteTmp()
        print("\n请输入问题或指令")
        usercmd = prompt(multiline=True,key_bindings=kb)
        Set.calc_functionnum()
        if usercmd == "help":
            print("输入deepthink打开深度思考\n输入websearch打开联网搜索\n输入uploadfile上传文件\n输入exit退出\n输入history打开历史记录\n输入new新建对话\n输入print输出对话历史\n其他输入开始对话")
            continue
        elif usercmd == "clear":
            os.system("clear")
        elif usercmd == "print":
            HistoryManager.PrintHistory()
        elif usercmd == "uploadfile":
            try:
                import aifile
                aifile.uploadfile()
                fileuploaded=True
            except:
                print("正在安装库，安装成功之后应用将关闭，请重新进入")
                os.system("bash install.sh")
                exit()
        elif usercmd=="deepthink":
            if Set.EnableDeepThink==False:
                Set.EnableDeepThink=True          
                print("深度思考已开启")
            else:
                Set.EnableDeepThink = False
                print("深度思考已关闭")  
        elif usercmd=="websearch":
            if Set.EnableSearch==False:
                Set.EnableSearch=True          
                print("联网搜索已开启")
            else:
                Set.EnableSearch = False
                print("联网搜索已关闭") 
        elif usercmd == "exit":
            HistoryManager.SaveDialog()
            os.remove("/home/shell/ai.config")
            exit()
        elif usercmd == "manageapps":
            pass
        elif usercmd=="generatepic":
            pass
        elif usercmd=="managemodels":
            ModelManager.SelectJsonFile()
        elif usercmd == "history":
            HistoryManager.LoadHistory()
            HistoryManager.PrintHistory()
        elif usercmd=="set":
            Set.SetDialog()
        elif usercmd == "new":
            HistoryManager.SaveDialog()
            HistoryManager.ChatHistory = []
            HistoryManager.IsNewChat = True
            HistoryManager.IsChanged = False
        elif usercmd!="":
            try:
                question = TextProcesser.CheckLen(HistoryManager.GetText(HistoryManager.ChatHistory, "user", usercmd))
                HistoryManager.IsChanged = True
                if Set.functionnum=="100":
                    print("当前是联网搜索模式")
                    print("Kimi-K2联网搜索:", flush=True)
                    assistant_response = PublicAI.GetWebAnswer(question)
                    HistoryManager.GetText(HistoryManager.ChatHistory, "assistant", assistant_response)
                elif Set.functionnum == "000":
                    print("当前是普通模式")
                    print("DeepSeek-V3.1：", flush=True)
                    assistant_response = PublicAI.GetAnswer(question)
                    HistoryManager.GetText(HistoryManager.ChatHistory, "assistant", assistant_response)
                elif Set.functionnum == "010":
                    print("当前是深度思考模式")
                    print("DeepSeek-V3.1:", flush=True)
                    assistant_response = PublicAI.GetDeepAnswer(question)
                    HistoryManager.GetText(HistoryManager.ChatHistory, "assistant", assistant_response)
                elif Set.functionnum== "110":
                    print("当前是深度思考+联网搜索")
                    print("Qwen-Plus", flush=True)
                    assistant_response = PublicAI.GetDeepWebAnswer(question)
                    HistoryManager.GetText(HistoryManager.ChatHistory, "assistant", assistant_response)
            except openai.RateLimitError:
                print("今日用量超过500，请明天再来")
            except TypeError:
                print("今日用量超过500或者涉及敏感信息")
            except UnicodeDecodeError:
                print("请不要输入特殊字符")
            except openai.AuthenticationError:
                print("API密钥无效或者超过1个月有效期")
